# -*- coding: utf-8 -*-
"""
Created on Mon Nov 22 17:26:50 2021

@author: kidwz
"""
import matplotlib.pyplot as plt
import numpy as np
import torch


# ======================================== #
# Functions related to 144K simulation
# ======================================== #
# data quant
# isint = 1 -> return quantized values as integer levels
# isint = 0 -> return quantized values as float numbers with the same range as input
def data_quantization_sym(data_float, half_level = 15, scale = None, isint = 0, clamp_std = None):
    if half_level <= 0:
        return data_float, 0

    std = data_float.std()
    if clamp_std != None and clamp_std != 0:
        data_float = torch.clamp(data_float, min = -clamp_std * std, max = clamp_std * std)

    if scale == None or scale == 0:
        scale = abs(data_float).max()
    if scale == 0:
        return data_float, 0

    data_quantized = (data_float / scale * half_level).round()

    if isint == 0:
        data_quantized = data_quantized * scale / half_level

    return data_quantized, scale


# 给 feature_map 加上 padding
# feature_map 维度： C，W，H
def feature_map_padding(feature_map, padding):
    while (len(feature_map.shape) < 3):
        feature_map = np.expand_dims(feature_map, axis = 0)
    feature_map_pad = np.pad(feature_map, ((0, 0), (padding, padding), (padding, padding)), mode = 'constant')
    return feature_map_pad


# 将忆阻器每层的输出 out_put 转换回 feature_map 的形式
# out_put shape = [W_out * H_out, C_out]
# feature_map shape = [C_out, W_out, H_out]
def output_to_feature_map(out_put, out_w, out_h):
    channels = out_put.shape[1]
    feature_map = out_put.transpose(1, 0).reshape([channels, out_w, out_h])
    return feature_map


# 将 feature_map 转化为下一层忆阻器的输入 array_input
# feature_map shape = [C_in, W_in, H_in]
# array_input shape = [W_out * H_out, C_out]
def feature_map_to_input(feature_map, kernel_size, stride, padding, repeat = None):
    while (len(feature_map.shape) < 3):
        feature_map = np.expand_dims(feature_map, axis = 0)
    in_channels = feature_map.shape[0]
    feature_in_w = feature_map.shape[1]
    feature_in_h = feature_map.shape[2]
    feature_out_w = int((feature_in_w - kernel_size + 2 * padding) / stride + 1)
    feature_out_h = int((feature_in_h - kernel_size + 2 * padding) / stride + 1)
    feature_map = feature_map_padding(feature_map, padding)
    input_rows = kernel_size ** 2 * in_channels
    output_rows = feature_out_w * feature_out_h
    array_input = np.zeros([input_rows, output_rows])
    idx = 0
    for i in range(feature_out_w):
        for j in range(feature_out_h):
            slide_window = feature_map[:, i * stride:i * stride + kernel_size,
                           j * stride:j * stride + kernel_size]
            array_input[:, idx] = slide_window.reshape(-1)
            idx += 1
    if repeat:
        array_input = np.tile(array_input, [repeat[0], 1])
    return array_input


# 如果权重复制，则返回一个复制的input
def input_repeat(input, row_repeat):
    input_rep = np.tile(input, [row_repeat, 1])
    return input_rep


# 池化
def pooling(feature_map, stride):
    channels = feature_map.shape[0]
    pooled_w = int(feature_map.shape[1] / stride)
    pooled_h = int(feature_map.shape[2] / stride)
    f2_pooling = np.zeros([channels, pooled_w, pooled_h])
    for i in range(channels):
        for row in range(pooled_w):
            for col in range(pooled_h):
                max_num = feature_map[i, row * 2:row * 2 + 2, col * 2:col * 2 + 2].max()
                f2_pooling[i, row, col] = max_num
    return f2_pooling


# 一列数据按脉冲展开
def col_bitwise_expansion(col_data, bitlen = None, return_bitlen = 0):
    col_data = col_data.reshape(-1, 1).astype(np.int32)
    col_data_ori = col_data + 0
    if bitlen == None:
        bitlen = int(abs(col_data).max())
    if bitlen == 0:
        return col_data
    rows = col_data.shape[0]
    col_data_expanded = []
    while ((col_data != 0).any()):
        bit_cur = (col_data > 0) * 1 + (col_data < 0) * -1
        col_data_expanded.append(bit_cur)
        col_data -= bit_cur
    col_data_expanded = np.concatenate(col_data_expanded, axis = 1).transpose(1, 0)
    np.random.shuffle(col_data_expanded)
    col_data_expanded = col_data_expanded.transpose(1, 0)
    assert (col_data_expanded.sum(axis = 1).reshape(-1, 1) == col_data_ori).all()
    return col_data_expanded


# feature数据按脉冲展开，每列都按照固定长度展开，即便该列都是0
# def feature_map_bitwise_expansion(feature_map, bitlen_max):
#     rows, cols = feature_map.shape
#     feature_map_quant, _ = data_quantization_sym(feature_map, isint = 1, half_level = bitlen_max)
#     feature_map_quant = feature_map_quant.astype(np.int32)
#     feature_map_expanded = np.zeros([rows, cols * int(bitlen_max)])
#     for i in range(cols):
#         feature_map_expanded[:, i * bitlen_max:(i + 1) * bitlen_max] = col_bitwise_expansion(
#             feature_map_quant[:, i],
#             bitlen_max)
#     return feature_map_expanded

# 根据每列输入的数据大小，自动做bitwise展开
# weight 和 input 可以是复制后的
# repeat 是权重复制的次数，是一个列表，[row_repeat, col_repeat]
def mvm_bitwise(input, weight, repeat = None):
    # bitwise乘加运算
    cal_times = input.shape[1]
    output_cols = weight.shape[1]
    output = np.zeros([cal_times, output_cols])
    for i in range(cal_times):
        col = input[:, i]
        bit_len = int(abs(col).max())
        col_expanded = col_bitwise_expansion(col, bit_len)
        temp = 0
        for j in range(bit_len):
            temp += (col_expanded[:, j].reshape(-1, 1) * weight).sum(axis = 0)
        output[i, :] = temp

    # 如果权重复制了，求出output的平均值
    if repeat:
        row_repeat = repeat[0]
        col_repeat = repeat[1]
        cols_output_avg = int(output_cols / col_repeat)
        output_avg = np.zeros([cal_times, cols_output_avg])
        for i in range(col_repeat):
            output_avg += output[:, i * cols_output_avg: (i + 1) * cols_output_avg]
        output_avg /= col_repeat
        output_avg /= row_repeat
        return output_avg
    return output


# 对二维数组进行并行bitwise展开，效率最高，无需循环
# input size = [rows, cols]
# 如果 dense != True:
# 返回值是一个展开后的稀疏二维数组，按照input中的最大值作为bitwise展开的位数，对每列进行展开
# output size = [rows, cols * bitlen]
# 如果 dense == True:
# 返回值是一个稠密二维数组，分别按照input中每列的最大值作为bitwise的展开位数。
# 即展开后都是0的列直接丢弃
# 在计算时需要知道bitlen_map中每个col展开了多少次，根据该值对144k的计算结果进行求和处理
def input_bitwise_expansion_fast(input, dense = True):
    # 如果输入为全0矩阵，直接返回其本身，在后续的mvm中跳过这次计算
    if (input == 0).all():
        return input, []

    if len(input.shape) == 1:
        input = input.reshape(-1, 1).astype(np.int32)
    input = input.astype(np.int32)
    input_ori = input + 0

    bitlen_map = np.max(abs(input), axis = 0)
    rows = input.shape[0]
    input_expanded = []

    while ((input != 0).any()):
        bit_cur = (input > 0) * 1 + (input < 0) * -1
        input_expanded.append(bit_cur.transpose(1, 0).reshape(-1, 1))
        input -= bit_cur

    input_expanded = np.concatenate(input_expanded, axis = 1)
    input_expanded = input_expanded.transpose(1, 0)
    np.random.shuffle(input_expanded)
    input_expanded = input_expanded.transpose(1, 0)
    input_expanded = np.split(input_expanded, int(input_expanded.shape[0] / rows), axis = 0)
    input_expanded = np.concatenate(input_expanded, axis = 1)

    if dense == True:
        mask = abs(input_expanded).sum(axis = 0)
        input_expanded = input_expanded[:, mask != 0]

    return input_expanded, bitlen_map


def mvm_bitwise_concat_push_fast(input, weight, repeat = None):
    # bitwise乘加运算
    cal_times = input.shape[1]
    output_cols = weight.shape[1]
    output = np.zeros([cal_times, output_cols])
    input_expanded, bitlen_map = input_bitwise_expansion_fast(input)

    # 如果feature map展开后的返回值为全0矩阵（即feature map）为全0，直接跳过这次运算，返回一个尺寸正确的全0矩阵作为结果
    if (input_expanded == 0).all():
        if repeat:
            col_repeat = repeat[1]
            return output[:, 0:int(output_cols / col_repeat)]
        else:
            return output
    output_bitwise = sdk_cal_sim(input_expanded, weight)

    output_bitwise_col = 0
    output_row = 0

    # 对计算结果按照展开的位数进行求和
    for j in bitlen_map:
        if j == 0:
            output[output_row] = 0
        else:
            output[output_row] = output_bitwise[output_bitwise_col:output_bitwise_col + j].sum(axis = 0)
        output_row += 1
        output_bitwise_col += j

    # 如果权重复制了，求出output的平均值
    if repeat:
        row_repeat = repeat[0]
        col_repeat = repeat[1]
        cols_output_avg = int(output_cols / col_repeat)
        output_avg = np.zeros([cal_times, cols_output_avg])
        for i in range(col_repeat):
            output_avg += output[:, i * cols_output_avg: (i + 1) * cols_output_avg]
        output_avg /= col_repeat
        output_avg /= row_repeat
        return output_avg
    return output


def mvm_bitwise_concat_push(input, weight, repeat = None):
    # bitwise乘加运算
    cal_times = input.shape[1]
    output_cols = weight.shape[1]
    output = np.zeros([cal_times, output_cols])
    bitlen_map = []
    input_expanded = []
    for i in range(cal_times):
        col = input[:, i]
        bit_len = int(abs(col).max())
        # log down bit length for current column
        bitlen_map.append(bit_len)
        # column bit-wise expansion
        col_expanded = col_bitwise_expansion(col, bit_len)
        input_expanded.append(col_expanded)
    input_expanded = np.concatenate(input_expanded, axis = 1)
    out_put_bitwise = sdk_cal_sim(input_expanded, weight)

    out_put_bitwise_col = 0
    output_row = 0
    for j in bitlen_map:
        if j == 0:
            j = 1
        output[output_row] = out_put_bitwise[out_put_bitwise_col:out_put_bitwise_col + j].sum(axis = 0)
        output_row += 1
        out_put_bitwise_col += j

    # 如果权重复制了，求出output的平均值
    if repeat:
        row_repeat = repeat[0]
        col_repeat = repeat[1]
        cols_output_avg = int(output_cols / col_repeat)
        output_avg = np.zeros([cal_times, cols_output_avg])
        for i in range(col_repeat):
            output_avg += output[:, i * cols_output_avg: (i + 1) * cols_output_avg]
        output_avg /= col_repeat
        output_avg /= row_repeat
        return output_avg
    return output


# simulate the mvm process in C200 SDK
def sdk_cal_sim(input, weight):
    cal_times = input.shape[1]
    sum = []
    for i in range(cal_times):
        temp = (input[:, i].reshape(-1, 1) * weight).sum(axis = 0)
        sum.append(temp.reshape(1, -1))
    result = np.concatenate(sum, axis = 0)
    return result.astype(np.int32)


# 重复的矩阵求出平均值
# weight是从忆阻器中读出的多个重复的矩阵
# repeat是这个矩阵重复的次数，格式为[row_repeat, col_repeat]
# 实际权重值的大小为 [weight.shape[0] / row_repeat, weight.shape[1] / col_repeat]
def weight_avg(weight, repeat):
    weight_avg = np.zeros([weight.shape[0] // repeat[0], weight.shape[1] // repeat[1]])
    m, n = weight_avg.shape
    for row in range(repeat[0]):
        for col in range(repeat[1]):
            weight_avg += weight[row * m:(row + 1) * m, col * n:(col + 1) * n]
    weight_avg /= (repeat[0] * repeat[1])
    return weight_avg


# 计算网络推理的softmax，即概率
def softmax(cross_bar_final_output):
    cross_bar_final_output = cross_bar_final_output.squeeze()
    max_num = cross_bar_final_output.max()
    cross_bar_final_output -= max_num
    prob = np.exp(cross_bar_final_output.max()) / np.exp(cross_bar_final_output).sum()
    return prob


# 显示图片
def show_img_and_result(col, result, label, prob):
    # plt.figure()
    plt.figure(figsize = (5, 7))
    plt.axis('off')
    plt.imshow(col.squeeze(), cmap = 'gray')
    if result == label:
        t = f'This is a: ' \
            f'\n {get_text_labels(result)} ' \
            f'\n probility = {prob:.2f} ' \
            f'\n CORRECT :)'
        plt.title(t, fontsize = 20, fontweight = 'bold', color = 'green')
    else:
        t = f'This is a: ' \
            f'\n {get_text_labels(result)} ' \
            f'\n probility = {prob:.2f} ' \
            f'\n WRONG!!'
        plt.title(t, fontsize = 20, fontweight = 'bold', color = 'red')
    plt.pause(0.1)


def get_text_labels(labels_num):
    text_labels = ['t-shirt', 'trouser', 'pullover', 'dress', 'coat', 'sandal', 'shirt', 'sneaker', 'bag',
                   'ankle boot']
    return text_labels[labels_num]
